import time

import coin.base.database.util as db_util
from concurrent.futures import ThreadPoolExecutor, as_completed
from coin.exchange.base.kr_rest.private_client_base import \
  ExWalletPrivateClientBase
from coin.blockchain.blockchair.btc.private_client import \
  BtcPrivateClient
from coin.blockchain.web3.eth.private_client import \
  EthPrivateClient
from coin.blockchain.web3.tron.private_client import \
  TronPrivateClient, TronWalletException
from coin.blockchain.cosmos.private_client import \
  CosmosPrivateClient
from coin.blockchain.finschia.private_client import \
  FinschiaPrivateClient
from coin.blockchain.web3.eth.constants import (
  ETH_TOKEN_TO_ADDRESS, BSC_TOKEN_TO_ADDRESS)
from coin.blockchain.web3.tron.constants import (
  TRON_TOKEN_TO_ADDRESS)
from coin.blockchain.cosmos.constants import (
  COSMOS_TOKEN_TO_ADDRESS, OSMOSIS_TOKEN_TO_ADDRESS)
from coin.blockchain.finschia.constants import (
  FNSA_TOKEN_TO_ADDRESS)
from coin.blockchain.request_util import parse_coin_key
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.proto.coin_query_pb2 import AccountBalance


API_KEY_CONFIG = '../../coin_deploy/support_accounting/blockchain_config/api_key_config.json'


def _get_private_client(blockchain_name):
  api_key_config = db_util.read_db_config(API_KEY_CONFIG)
  if blockchain_name == 'bitcoin':
    return BtcPrivateClient(blockchain_name=blockchain_name,
                            api_key=api_key_config.get(blockchain_name))
  elif blockchain_name == 'tron':
    return TronPrivateClient(blockchain_name=blockchain_name,
                             api_key=api_key_config.get(blockchain_name))
  elif blockchain_name in ('ethereum', 'bsc'):
    return EthPrivateClient(blockchain_name=blockchain_name,
                            api_key=api_key_config.get(blockchain_name))
  elif blockchain_name in ('cosmos', 'osmosis'):
    return CosmosPrivateClient(blockchain_name=blockchain_name,
                               api_key=api_key_config.get(blockchain_name))
  elif blockchain_name == 'finschia':
    return FinschiaPrivateClient(blockchain_name=blockchain_name,
                                 api_key=api_key_config.get(blockchain_name))
  else:
    raise ValueError(f'{blockchain_name} blockchain not implemented')


def _get_supported_currency_list(blockchain_name):
  if blockchain_name == 'bitcoin':
    return ['BTC']
  elif blockchain_name == 'ethereum':
    return ETH_TOKEN_TO_ADDRESS.keys()
  elif blockchain_name == 'bsc':
    return BSC_TOKEN_TO_ADDRESS.keys()
  elif blockchain_name == 'tron':
    return TRON_TOKEN_TO_ADDRESS.keys()
  elif blockchain_name == 'cosmos':
    return COSMOS_TOKEN_TO_ADDRESS.keys()
  elif blockchain_name == 'osmosis':
    return OSMOSIS_TOKEN_TO_ADDRESS.keys()
  elif blockchain_name == 'finschia':
    return FNSA_TOKEN_TO_ADDRESS.keys()
  else:
    raise ValueError(f'{blockchain_name} blockchain not implemented')


def _generate_account_balance_update(account_key, req_timestamp, balance_list):
  account_balance = AccountBalance(
      market_type = account_key.market_type or 'ExWallet',
      exchange = account_key.exchange or 'HwWallet',
      owner = account_key.owner,
      query_ts = int(req_timestamp * 10 ** 9),
      each_balance = balance_list
  )
  return RestUpdate(req_timestamp, time.time(), account_balance)


class HwWalletPrivateClient(ExWalletPrivateClientBase):
  def __init__(self, key_file):
    self._key = parse_coin_key(key_file)
    self._blockchains = self._key.chain_keys

  def query_account_balance_impl(self):
    req_timestamp = time.time()
    with ThreadPoolExecutor(max_workers=24) as executor:
      futures_list = []
      balance_dict = {}
      for blockchain in self._key.chain_keys:
        blockchain_name = blockchain.blockchain
        address = blockchain.access_key
        privc = _get_private_client(blockchain_name)
        currency_list = _get_supported_currency_list(blockchain_name)
        for currency in currency_list:
          f = executor.submit(privc.query_account_balance_impl,
                              address=address,
                              currency=currency)
          futures_list.append(f)
      for f in as_completed(futures_list):
        try:
          balance = f.result()
          if balance.currency not in balance_dict:
            balance_dict[balance.currency] = balance
          else:
            balance_dict[balance.currency].total += balance.total
        except TronWalletException as e:
          print(type(e), e)

    update = _generate_account_balance_update(
      self._key, req_timestamp, balance_dict.values())
    return update

  def get_latest_block_number_timestamp(self):
    return self._privc.get_latest_block_number_timestamp()

  def get_block_timestamp(self, block_number):
    return self._privc.get_block_timestamp(block_number)


if __name__ == '__main__':
  hw_wallet_client = HwWalletPrivateClient(
    key_file='../../coin_key/view/ledger04/hw_wallet/view_key.json')
  print(hw_wallet_client.query_account_balance())
